# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: taekwon

import datetime
import gzip
import logging
import os
import pickle
import typing

from recordclass import dataobject

from coin.base.datetime_util import to_datetime
from coin.strategy.mm.subscription import FeedSubscriptionRequest
from coin.exchange.base.tick import FixedTickPrice

import experimental.prophet as prophet
from coin.experimental.jaewon.research.adaptive_edge import features, lt_model
from experimental.prophet.ops.linear_model import linear_regression_predict
from experimental.prophet.craps.universe import universe_all


def _model(symbol_nick, ref_symbol_nicks, linear_model, alpha_model, date, debug=False):
  # Time
  ts = prophet.global_timestamp()

  # Symbols
  universe = {elem.nickname: elem for elem in universe_all(date)}
  symbol_info = universe[symbol_nick]
  target_coin = symbol_info.product.base.currency.lower()
  assert target_coin in ('btc', 'eth', 'eos')

  feed = prophet.fastfeed_coin(symbol_info.product, symbol_info.sub_req)

  book = prophet.book(feed, inverse_product=symbol_info.inverse)
  trade = prophet.create_node('fastfeed.Trade', [feed.trade]).outputs
  ask0_p = book.estimated_ask0_price.named_ref('ask0_p')
  bid0_p = book.estimated_bid0_price.named_ref('bid0_p')

  # Features
  with prophet.namespace('feature'):
    feature_vars = features.features(target_coin, ref_symbol_nicks + [symbol_nick], date)

  # TODO(taekwon): do this later.
  # Sampler
  # with prophet.namespace('sampler'):
  #   sampler_vars = features.samplers(target_coin, 'feature/')
  #   sampler_vars = {v.name: v for v in sampler_vars}
  # Trigger
  # trigger = features.trigger(target_coin, symbol_nick, 'feature/')  

  if debug:
    ts_10s = prophet.timer('10s')
    with prophet.control_if(ts % (10*10**9) == 0):
      prophet.printer('%t %s %s', ts, buy_p, sell_p)
  
  return feature_vars


class PassModelConfig(dataobject):
  tick: typing.Any = None
  machine: str = ''
  symbol_nick: str
  ref_symbol_nicks: typing.List[str] = []
  linear_model_filepath: str
  alpha_model: typeing.Any = None
  sampler_offsets: typing.Any = None
  pull_edge_bp: float = 2.
  bound_price: bool = True
  lt_model_param: typing.Any = None
  ref_time: typing.Any = None


class PassModel:
  def __init__(self, config, pos_getter=None, trigger_callback=None):
    self._config = config
    self._pos_getter = pos_getter
    self.trigger_callback = trigger_callback

    if self._config.ref_time is not None:
      self._ref_time = to_datetime(self._config.ref_time)
    else:
      self._ref_time = datetime.datetime.now()

    self._graph = None
    self._adapter = None
    self._init_graph()

    self._tick = None
    self._init_tick()

    self._vars = []
    self._var_names = []
    self._rows = []
    self._lookback = 10

    self._ask0_p_var = self._graph.get_named_variable('ask0_p')
    self._bid0_p_var = self._graph.get_named_variable('bid0_p')
    self._buy_p_var = self._graph.get_named_variable('edge_model/buy_p')
    self._sell_p_var = self._graph.get_named_variable('edge_model/sell_p')
    self._timestamp = self._graph.get_named_variable('global_timestamp')
    self._trigger = self._graph.get_named_variable('trigger')

    self._buy_price = None
    self._buy_price_pull = None
    self._sell_price = None
    self._sell_price_pull = None

    self._last_triggered = 0

  def _init_tick(self):
    if self._config.tick is not None:
      tick = self._config.tick
      self._tick = FixedTickPrice(tick) if isinstance(tick, float) else tick
      return

    universe = {elem.nickname: elem for elem in universe_all(self._ref_time)}
    symbol_info = universe[self._config.symbol_nick]
    self._tick = symbol_info.product_info.gen_tick_price()
    assert isinstance(self._tick, FixedTickPrice)

  def _init_graph(self):
    # TODO(taekwon): ridge load
    linear_model_filepath = os.path.expanduser(self._config.linear_model_filepath)
    with gzip.open(linear_model_filepath, 'rb') as f:
      linear_model = pickle.load(f)

    self._graph = prophet.Graph()
    with prophet.as_default_graph(self._graph):
      _model(self._config.symbol_nick, self._config.ref_symbol_nicks, linear_model, 
             self._config.alpha_model, self._ref_time)

    self._adapter = prophet.FastfeedGraphAdapter(
        self._graph, self._config.machine, eval_callback=self._on_eval)

  def _on_eval(self, graph, live=False):
    if not self.ready:
      return

    ts = self.timestamp
    if not self.triggered and ts - self._last_triggered < 10*(10**9):
      return
    self._last_triggered = ts

    # TODO(taekwon):
    # self._buy_price_pull = self._buy_price * (1. + self._config.pull_edge_bp / 10000.)
    # self._sell_price_pull = self._sell_price * (1. - self._config.pull_edge_bp / 10000.)
    
    # Model    
    x_vars = [v for v in feature_vars if not v.name.endswith('.x')]
    row = [var.value for var in self._vars]
    self._rows = self._rows[-self._lookback+1:]    
    self._rows.append(row)

    df = pd.DataFrame(self._rows, columns=self._var_names)
    df = make_df(x_vars)
    df = 

    position = linear_model['intercept']
    for alpha in alphas:
      position += linear_model[alpha.name] * alpha.func(df)
    
    position.named_ref('position')

    if self.trigger_callback:
      self.trigger_callback()

  def prepare(self, feed_subsystem):
    feed_subsystem.subscribe_fastfeed(
        self._adapter.get_feed_sub_requests(),
        self._adapter.on_book_reset)

  @property
  def ready(self):
    return not self._buy_p_var.is_null and not self._sell_p_var.is_null

  @property
  def triggered(self):
    return self._trigger.value

  @property
  def timestamp(self):
    return self._timestamp.value

  @property
  def buy_price(self):
    return self._buy_price

  @property
  def sell_price(self):
    return self._sell_price

  @property
  def buy_price_pull(self):
    return self._buy_price_pull

  @property
  def sell_price_pull(self):
    return self._sell_price_pull
